import torch
import torch.nn as nn
from torchvision.models import mobilenet_v2

class MobileNetV2Classifier(nn.Module):
    def __init__(self, num_classes=7, width_mult=0.5):
        super(MobileNetV2Classifier, self).__init__()
        
        # 创建MobileNetV2骨干网络，不加载预训练权重
        self.backbone = mobilenet_v2(pretrained=False, width_mult=width_mult)
        
        # 获取最后一层通道数
        last_channel = self.backbone.last_channel
        
        # 分类器
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(last_channel, num_classes)
        )
    
    def forward(self, x):
        x = self.backbone.features(x)
        x = self.classifier(x)
        return x

def create_model(num_classes=7, width_mult=0.5):
    """创建MobileNetV2分类模型"""
    return MobileNetV2Classifier(num_classes, width_mult)

if __name__ == "__main__":
    # 简单测试
    model = create_model()
    x = torch.randn(2, 3, 224, 224)
    output = model(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"模型参数量: {sum(p.numel() for p in model.parameters()):,}")